# -*- coding: utf-8 -*-

import numpy as np
import segyio
from numpy import array
from scipy.interpolate import griddata


class Seismic(object):

    def __init__(self, data, dtype=float, params=None):

        if params is None:
            params = {}

        self.params = params
        # np.array：由于总是创建副本，因此当你需要确保数据不会被原始数据修改时，使用 np.array 更为安全
        # np.asarray：当你希望尽可能避免不必要的数据复制时，使用 np.asarray，特别是在处理大型数据集时，这可以节省内存和时间
        self.data = np.asarray(data, dtype=dtype)  # 将其他类型的数据转换为numpy数组
        self.ntraces = params.get('ntraces', self.data.shape[-2])
        self.dt = params.get('dt', 0.001)
        self.cdp = params.get('cdp')

        self.fb = params.get('fb')
        # shape返回一个表示数组各维度大小的元组
        self.nsamples = int(self.data.shape[-1])  # 多维数组的最后一个维度的大小
        self.data = np.squeeze(self.data)  # 去除数组中所有长度为1的维度
        self.cdpx = params.get('cdpx')
        self.cdpy = params.get('cdpy')
        self.wellx = params.get('wellx')
        self.welly = params.get('welly')
        self.inlin = params.get('inlin')
        self.crosslin = params.get('crosslin')
        self.hcf = params.get('hcf')
        # 修改时间长度
        # 生成一个等差数列，start为0，end为nsamples*dt（不包含 end），step为dt
        self.tbasis = np.arange(0, self.nsamples * self.dt, self.dt)
        return

    @classmethod
    # cls: 类本身，segy_file: segy文件路径，params: 参数字典
    def from_segyio(cls, segy_file, params=None):
        params = {}
        with segyio.open(segy_file, strict=False, ignore_geometry=True) as s:
            # Read the data.将每个trace的数据堆叠起来，形成三维数组
            # 将t.astype中所有数据的数据类型转换为float
            data = np.stack([t.astype(np.float) for t in s.trace])
            # Get the trace numbers.
            # segyio.TraceField.CDP获取CDP字段的索引
            # “道集号”，它代表了一组地震道的集合，这些地震道共同记录了从不同震源点激发的地震波在地下某个深度点的反射
            cdp = np.array([t[segyio.TraceField.CDP] for t in s.header])
            # 开始消噪时间，有助于在地震数据的预处理阶段去除噪声，提高数据质量
            fb = np.array([t[segyio.TraceField.MuteTimeStart]
                          for t in s.header])
            # CDP的X坐标或者与CDP相关的某个横向线号（Cross-line）的X坐标
            # 不同震源点激发的地震波会在地下某个深度点产生反射，这个深度点就是CDP
            # CDP技术允许地震学家通过多个不同震源点激发的地震波在地下的同一个点反射回来的信号，来提高数据的准确性和分辨率
            cdpx = np.array([t[segyio.TraceField.CDP_X] for t in s.header])
            # 公共深度点CDP相关的Y坐标
            cdpy = np.array([t[segyio.TraceField.CDP_Y] for t in s.header])
            # 接收器组的X坐标，有助于确定地震道在三维空间中的水平位置，和GroupY 共同定义了地震道在地面上的地理位置
            # Group通常指的是接收器组，即一组地震检波器的集合
            # 用于接收从地下反射回来的地震波
            # Group数据有助于确定地震波的到达时间和振幅，从而可以计算出地震波的传播时间和路径，进而推断地下结构的特性
            wellx = np.array([t[segyio.TraceField.GroupX] for t in s.header])
            welly = np.array([t[segyio.TraceField.GroupY] for t in s.header])
            # INLINE_3D 字段通常用于标识每个地震道属于哪个测线
            # 在三维地震数据中，每个地震道都会被分配一个横线和纵线的编号，用于标识其在三维空间中的位置
            # 在三维地震勘探中，数据通常按照测线（inline）和横线（crossline）进行组织，形成了一个三维的数据体
            inlin = np.array([t[segyio.TraceField.INLINE_3D]
                             for t in s.header])
            crosslin = np.array([t[segyio.TraceField.CROSSLINE_3D]
                                for t in s.header])
            # 地震数据在处理过程中应用的高频截止值，以便于了解数据处理的历史和参数
            # 高通滤波器用于去除地震数据中的低频成分，只允许高于截止频率的信号通过
            # 这在地震数据处理中是一个常见的步骤，用于减少噪声和提高数据的分辨率
            hcf = np.array([t[segyio.TraceField.HighCutFrequency]
                           for t in s.header])

            params['cdpx'] = cdpx  # CDP的X坐标
            params['cdpy'] = cdpy  # CDP的Y坐标
            params['wellx'] = wellx  # 接收器组的X坐标
            params['welly'] = welly  # 接收器组的Y坐标
            params['cdp'] = cdp  # CDP
            params['fb'] = fb  # 开始消噪时间
            params['inlin'] = inlin  # 地震道横线编号
            params['crosslin'] = crosslin  # 地震道纵线编号
            params['hcf'] = hcf  # 高频截止值
            # 采样间隔，可以影响到数据的采样率和频率带宽
            params['dt'] = s.bin[segyio.BinField.Interval] / 1000

        return cls(data, dtype=float, params=params)

    @classmethod
    def from_segy_with_segyio(cls, segy_file, params=None):
        return cls.from_segyio(segy_file, params=params)

    from_segy = from_segy_with_segyio

    def get_data(self,
                 l=1):
        return self.data

    def imshow(self,
               l=1,
               ax=None,
               gain=1.0,
               ts=0.0,
               td=6.0,
               lped=0.0,
               fcx=0,
               lpsd=0.0,
               dep=[],
               kb=0.0,
               secai='seismic',
               ):
        data = self.get_data(l)
        stp = int(ts * 1000.0 / self.dt)
        etp = int(td * 1000.0 / self.dt)
        fbmax = 0.0
        fbmin = 8000.0

        plotdata = []
        plotdep = []
        dep1 = np.array(self.cdpx)
        for brfjj, x, trace in zip(range(0, len(dep1)), dep1, data):
            if lpsd + kb <= x <= lped + kb:
                plotdata.append(trace[stp:etp])
                plotdep.append(x)
                if self.fb[brfjj] > fbmax:
                    fbmax = self.fb[brfjj]
                if self.fb[brfjj] < fbmin:
                    fbmin = self.fb[brfjj]
        if fbmin < ts * 1000.0:
            fbmin = ts * 1000.0
        if fbmax > td * 1000.0:
            fbmax = td * 1000.0

        brfpp = 0
        fctime = np.arange(len(dep))
        if fcx == 1:
            # fctime = np.arange(len(dep))
            for ii in dep:  # ii为深度
                for jj in range(1, len(self.cdpx)):
                    if ii == self.cdpx[jj] + kb:
                        fctime[brfpp] = self.fb[jj]
                        brfpp = brfpp + 1
                    elif ((self.cdpx[jj - 1] + kb < ii < self.cdpx[jj] + kb)
                          or
                            (self.cdpx[jj - 1] + kb > ii > self.cdpx[jj] + kb)
                          ):
                        fctime[brfpp] = (ii - self.cdpx[jj]) * (self.fb[jj - 1] - self.fb[jj]) / (self.cdpx[jj - 1] - self.cdpx[jj]) + self.fb[jj]
                        brfpp = brfpp + 1

        time = np.arange(ts * 1000, td * 1000, self.dt)
        py, px = np.meshgrid(plotdep, time)
        points = np.c_[py.ravel(), px.ravel()]
        grid_x, grid_y = np.mgrid[lpsd:lped:10,
                                  ts * 1000.0:td * 1000.0:self.dt]

        plotdata = array(plotdata)
        data2 = plotdata.flatten('F')

        grid_z2 = griddata(points, data2, (grid_x, grid_y), method='nearest')
        data1 = gain * grid_z2.transpose()

        ax.imshow(data1, cmap=secai, aspect='auto', vmin=-0.5,
                  vmax=0.5, extent=[lpsd, lped, td * 1000, ts * 1000],
                  interpolation='bilinear')

        return [fbmin, fbmax, ax, fctime]

    def imshow2(self,
                l=1,
                ax=None,
                gain=1.0,
                ts=0.0,
                td=6.0,
                lped=0.0,
                lpsd=0.0,
                secai='seismic',
                ):
        data = self.get_data(l)
        stp = int(ts * 1000.0 / self.dt)
        etp = int(td * 1000.0 / self.dt)

        plotdata = []
        plotdep = []
        dep1 = np.array(self.cdpx)
        for brfjj, x, trace in zip(range(0, len(dep1)), dep1, data):
            if lpsd <= x <= lped:
                plotdata.append(trace[stp:etp])
                plotdep.append(x)

        time = np.arange(ts * 1000, td * 1000, self.dt)
        py, px = np.meshgrid(plotdep, time)
        points = np.c_[py.ravel(), px.ravel()]
        grid_x, grid_y = np.mgrid[lpsd:lped:10,
                                  ts * 1000.0:td * 1000.0:self.dt]

        plotdata = array(plotdata)
        data2 = plotdata.flatten('F')

        grid_z2 = griddata(points, data2, (grid_x, grid_y), method='nearest')
        data1 = gain * grid_z2.transpose()

        ax.imshow(data1, cmap=secai, aspect='auto', vmin=-0.5, vmax=0.5,
                  extent=[lpsd, lped, td * 1000, ts * 1000],
                  interpolation='bilinear')
        return ax

    def wiggle_plot2(self,
                     l=1,
                     ax=None,
                     perc=99.0,
                     gain=1.0,
                     rgb=(0, 0, 0),
                     alpha=0.5,
                     lw=0.2,
                     ts=0.0,
                     td=6.0,
                     times=0,
                     stt=0,
                     sein=0,
                     fcx=0,
                     zu=[],
                     py=[],
                     fctime=[],
                     font=[],
                     fbmin=0,
                     fbmax=0,
                     sttflag=0,
                     plotzflag=0,
                     ):

        data = self.get_data(l)
        rgba = list(rgb) + [alpha]
        sc = np.percentile(data, perc)  # Normalization factor
        stp = int((ts * 1000.0 + times) / self.dt)
        etp = int((td * 1000.0 + times) / self.dt)
        if sttflag == 0:
            wigdata = data[stt:stt + sein, stp: etp]
        else:
            wigdata = data[stt - sein:stt, stp: etp]
        xpos = np.arange(20, 20 * sein + 20, 20)

        for x, trace in zip(xpos, wigdata):
            # Compute high resolution trace.
            amp = x + gain * 20 * trace / sc

            t = np.arange(stp * self.dt, etp * self.dt, self.dt)
            hypertime = np.linspace(t[0], t[-1], (10 * t.size - 1) + 1)
            hyperamp = np.interp(hypertime, t, amp)

            ax.plot(hyperamp, hypertime, 'k', lw=1.2 * lw)
            ax.fill_betweenx(hypertime, hyperamp, x,
                             where=hyperamp > x,
                             facecolor=rgba,
                             lw=0,
                             )
        if fcx == 1:
            for ii in range(0, len(fctime)):
                if fctime[ii] > fbmax or fctime[ii] < fbmin:
                    continue
                if sttflag == 1:
                    ax.plot([20 * sein - 180, 21 * sein], [fctime[ii] +
                            times, fctime[ii] + times], 'r', lw=2.5 * lw)
                else:
                    ax.plot([0, 200], [fctime[ii] + times,
                            fctime[ii] + times], 'r', lw=2.5 * lw)
                if plotzflag == 1:
                    ax.text(200, fctime[ii] + times + py[ii], zu[ii],
                            transform=ax.transData,
                            verticalalignment='center',
                            horizontalalignment='left', fontsize=font[ii],
                            color='r')

        ax.set_xlim(0, 20 * sein + 20)
        ax.invert_yaxis()
        # ax.invert_xaxis()

        return ax
